package com.coscon.filter;


import com.coscon.utils.XSSUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.ArrayUtils;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.util.HashMap;
import java.util.Map;

/**
* XSS过滤
* @author 409390047@qq.com
* @date 2021/3/12
 */
@Slf4j
public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {
    HttpServletRequest orgRequest = null;

    
    //排除的input name数组
    private String[] extInputNames;
    

    public XssHttpServletRequestWrapper(HttpServletRequest request) {
        super(request);
        orgRequest = request;
    }

    /**
     * 覆盖getParameter方法，将参数名和参数值都做xss过滤。<br/>
     * 如果需要获得原始的值，则通过super.getParameterValues(name)来获取<br/>
     * getParameterNames,getParameterValues和getParameterMap也可能需要覆盖
     * 不转义的办法,直接返回 return super.getParameter(xssEncode(name));
     */
    @Override
    public String getParameter(String name) {
    	if(this.getExtInputNames()!=null && ArrayUtils.contains(this.getExtInputNames(), name)){
    		return super.getParameter(xssEncodeParameter(name));//排除的，则不转义
    	}
        String value = super.getParameter(xssEncodeParameter(name));

        if (value != null) {
            value = xssEncodeParameter(value);
        }
        return value;
    }
    @Override
    public String getQueryString() {
        return super.getQueryString() == null ? "" : xssEncodeParameter(super.getQueryString()).toString();
    }
    @Override
	public String[] getParameterValues(String name) {
    	if(this.getExtInputNames()!=null && ArrayUtils.contains(this.getExtInputNames(), name)){
    		return super.getParameterValues(xssEncodeParameter(name));//排除的，则不转义
    	}
    	String[] array = super.getParameterValues(name);
    	if(array!=null && array.length>0){
    		for(int i = 0; i < array.length;i++){
    			array[i] = xssEncodeParameter(array[i]);
    		}
    	}
		return array;
	}
    
	 /**
     * 重写并过滤getParameterMap方法
     */
    @Override
    public Map<String,String[]> getParameterMap() {
		Map<String,String[]> paraMap = super.getParameterMap();
        // 对于paraMap为空的直接return
        if (null == paraMap || paraMap.isEmpty()) {
            return paraMap;
        }
        //super.getParameterMap()不允许任何修改，所以只能做深拷贝
        Map<String, String[]> paraMapCopy = new HashMap<String, String[]>();
        //实际上putAll只对基本类型深拷贝有效，如果是自定义类型，则要找其他办法
        paraMapCopy.putAll(paraMap);
        
        for (Map.Entry<String, String[]> entry : paraMapCopy.entrySet()) {
        	if(this.getExtInputNames()==null || !ArrayUtils.contains(this.getExtInputNames(), entry.getKey()) ){
        		//进入编码处理，
        		String[] values     = entry.getValue();
        		if (null == values) {
        			continue;
        		}
        		String[] newValues  = new String[values.length];
        		for (int i = 0; i < values.length; i++) {
        			newValues[i] = xssEncodeParameter(values[i]);
        		}
        		entry.setValue(newValues);
        	}//如果处于排除的，则不处理
        }
        return paraMapCopy;
    }
	
	
	/**
     * 覆盖getHeader方法，将参数名和参数值都做xss过滤。<br/>
     * 如果需要获得原始的值，则通过super.getHeaders(name)来获取<br/>
     * getHeaderNames 也可能需要覆盖
     */
    @Override
    public String getHeader(String name) {
        String value = super.getHeader(xssEncodeParameter(name));
        if (value != null) {
            value = XSSUtils.xssEncodeParameter(value);
        }
        return value;
    }
    /**
     * 将容易引起xss漏洞的半角字符直接替换成全角字符
     * 对表单或者参数
     * @param s
     * @return
     */
    private  String xssEncodeParameter(String s) {
        if (s == null || s.isEmpty()) {
            return s;
        }
        return XSSUtils.xssEncodeParameter(s);
    }
    
    

    /**
     * 获取最原始的request
     *
     * @return
     */
    public HttpServletRequest getOrgRequest() {
        return orgRequest;
    }
    
    
	public String[] getExtInputNames() {
        if(extInputNames!=null){
            return (String[]) extInputNames.clone();
        }else{
            return null;
        }

	}
	public void setExtInputNames(String[] extInputNames) {
        if(extInputNames!=null){
            this.extInputNames = extInputNames.clone();
        }
	}
    
    
    
    
    /**
     * 获取最原始的request的方法
     * @return
     */
    public HttpServletRequest getOrgRequest(HttpServletRequest req) {
        if (req instanceof XssHttpServletRequestWrapper) {
            return ((XssHttpServletRequestWrapper) req).getOrgRequest();
        }
        return req;
    }
}